# This code is part of LINKEQ.
#
# (C) Copyright LINKE 2023.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
#
# -*- coding: utf-8 -*-
# @Time     : 2023/6/7 19:40
# @Author   : HFALSH @ LINKE
# @File     : simulator.py
# @Software : PyCharm
import numpy as np
import qiskit
from qiskit.circuit.random import random_circuit
from quantumcircuit.circuit import QuantumCircuit
from quantumcircuit.gate import *
from backend import Backend
from backend.simulator_cpp import *

# using qulacs simulator
# [1] Suzuki Y, Kawase Y, Masumura Y, et al. Qulacs: a fast and versatile quantum circuit simulator for research purpose[J]. Quantum, 2021, 5: 559.
import qulacs
from qulacs import QuantumState, PauliOperator, DensityMatrix
from qulacs.converter import convert_QASM_to_qulacs_circuit, convert_qulacs_circuit_to_QASM

from collections import Counter

class Simulator(Backend):
    def __init__(self, name="Qulacs simulator", capacity=None, info=None):
        super().__init__(name)
        self.name = name
        self.n_qubits = capacity
        self.info = info
        self.initila_state = None

    def get_backend_info(self):
        return f"Backend: {self.name}, Capacity: {self.capacity}, Info:{self.info}"

    def set_initial_state(self, mat:np.ndarray):
        '''
        根据矩阵mat生成初始量子态
        '''
        dim = mat.shape
        try:
            if not (dim[0] == dim[1]):
                raise ValueError("input array's shape is not square!")
        except ValueError as e:
            print("raise exception:", repr(e))

        state = DensityMatrix(int(np.log2(dim[0])))
        state.load(mat)
        self.initila_state = state


    def run(self, qcircuit, *args):
        '''
        todolist:
            1. 提供QuantumCircuit到QASM的转化功能；
            2. 设计result对象；
        '''
        qulacs_circuit = qcircuit
        if type(qcircuit) == QuantumCircuit and self.name == "Qulacs simulator":
            qasm_list = qcircuit.to_QASM_list()
            qulacs_circuit = convert_QASM_to_qulacs_circuit(qasm_list)
            pass
        elif type(qcircuit) == QuantumCircuit and self.name == "C++ simulator":
            sim_cpp = QuantumSimulator()
            state = sim_cpp.simulate(qcircuit.qubit_number, qcircuit.to_Operation_list())
            return state
        elif type(qcircuit) == qulacs.QuantumCircuit:
            qulacs_circuit = qcircuit
        if not self.initila_state:
            state = QuantumState(qulacs_circuit.get_qubit_count())
        else:
            state = self.initila_state
        qulacs_circuit.update_quantum_state(state)
        self.result_state = state
        return state
        # pass

    def state_pauli_expectation(self, pauli_str):
        expectation_value = 0
        for i in range(len(pauli_str)):
            for key, value in pauli_str[i].items():
                pauli = PauliOperator(pauli_string=key, coef=value)
                value = pauli.get_expectation_value(self.result_state).real
                # print("expect", value)
                expectation_value += value
        return expectation_value

    def sample_most_likely(self, shots: int = 1024):
        data = self.result_state.sampling(shots)
        sample_data = [format(value, "b").zfill(self.result_state.get_qubit_count()) for value in data]
        counts = Counter(sample_data)
        most_common = counts.most_common(1)
        self.sample_data = counts
        return most_common[0][0]

import time
import matplotlib.pyplot as plt

def create_circuit(n_qubits, depth):
    """创建量子电路"""
    qc = QuantumCircuit(n_qubits)
    
    for d in range(depth):
        # 第一层：H和RX门
        for i in range(n_qubits):
            qc.add_gate(H(i))
            qc.add_gate(RX(i,0.5))
        
        # 第二层：CX门
        for i in range(n_qubits-1):
            qc.add_gate(CX(i,i+1))
        
        # 第三层：X和RX门
        for i in range(n_qubits):
            qc.add_gate(X(i))
            qc.add_gate(RX(i,0.5))
        
        # 第四层：CX门
        for i in range(n_qubits-1):
            qc.add_gate(CX(i,i+1))
    
    return qc

def run_test(n_qubits_list, depth_list):
    """运行测试并记录时间"""
    results = []
    
    for n in n_qubits_list:
        for d in depth_list:
            print(f"Testing n_qubits={n}, depth={d}")
            
            # 创建电路
            qc = create_circuit(n, d)
            
            # 测试Python模拟器
            start_time = time.time()
            sim1 = Simulator()
            st1 = sim1.run(qc)
            time1 = time.time() - start_time
            
            # 测试C++模拟器
            start_time = time.time()
            sim2 = Simulator(name="C++ simulator")
            st2 = sim2.run(qc)
            time2 = time.time() - start_time
            
            results.append((n, d, time1, time2))
            
            # 写入文件
            with open("result.txt", "a") as f:
                f.write(f"{n}\t{d}\t{time1:.4f}\t{time2:.4f}\n")
    
    return results

def plot_results(results):
    """绘制结果图表"""
    # 按比特数分组
    n_qubits = sorted(set(r[0] for r in results))
    depths = sorted(set(r[1] for r in results))
    
    # 准备数据
    py_times = [[] for _ in n_qubits]
    cpp_times = [[] for _ in n_qubits]
    
    for n, d, t1, t2 in results:
        idx = n_qubits.index(n)
        py_times[idx].append(t1)
        cpp_times[idx].append(t2)
    
    # 创建图表
    plt.figure(figsize=(12, 6))
    
    for i, n in enumerate(n_qubits):
        plt.subplot(1, 2, 1)
        plt.plot(depths, py_times[i], 'o-', label=f'n={n} (Python)')
        plt.subplot(1, 2, 2)
        plt.plot(depths, cpp_times[i], 's-', label=f'n={n} (C++)')
    
    # 设置图表属性
    for i in range(2):
        plt.subplot(1, 2, i+1)
        plt.xlabel('Circuit Depth')
        plt.ylabel('Time (s)')
        plt.title('Python Simulator' if i == 0 else 'C++ Simulator')
        plt.legend()
        plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('output.png')
    plt.show()

def test():
    # 测试参数
    n_qubits_list = [28]
    depth_list = [4]
    
    # 清空结果文件
    with open("result.txt", "w") as f:
        f.write("n_qubits\tdepth\tpython_time\tcpp_time\n")
    
    # 运行测试
    results = run_test(n_qubits_list, depth_list)
    
    # 绘制图表
    plot_results(results)
if __name__ == '__main__':
    # qiskit_qc = random_circuit(3, 10)
    # # qc = qulacs.conve
    # qulacs_qc = convert_QASM_to_qulacs_circuit(qiskit_qc)
    # print(qulacs_qc)

    # n = 10
    # circuit = qulacs.QuantumCircuit(n)
    # circuit.add_H_gate(0)
    # circuit.add_X_gate(2)
    #
    # qasm_list = convert_qulacs_circuit_to_QASM(circuit)
    #
    # print(qasm_list)
    #
    # qulacs_qc = convert_QASM_to_qulacs_circuit(qasm_list)
    #
    # print(qulacs_qc)
    #
    # sim = Simulator()
    # result = sim.run(qcircuit=qulacs_qc)
    # # print(result)
    #
    # pauli_str = [{'Z 0 Z 1': 1}, {'Z 0 Z 3': 1}, {'Z 0 Z 9': 1}, {'Z 1 Z 2': 1}, {'Z 1 Z 3': 1}, {'Z 3 Z 8': 1},
    #              {'Z 8 Z 6': 1}, {'Z 8 Z 4': 1}, {'Z 2 Z 7': 1}, {'Z 2 Z 9': 1}, {'Z 7 Z 5': 1}, {'Z 7 Z 4': 1},
    #              {'Z 6 Z 4': 1}, {'Z 6 Z 5': 1}, {'Z 9 Z 5': 1}]
    # expec = sim.state_pauli_expectation(pauli_str=pauli_str)
    # print(expec)
    # sample_data = sim.sample_most_likely()
    # print(sample_data)

    test()




